#include "mpu6050_light_dmp.h"
#include "Driver/CLOCK/clock.h"
#include "Driver/hardware/HMC5883L/HMC5883L.h"

/****************************************************************************************************************/
/*                                        私有变量                                                             */
/****************************************************************************************************************/

static MPU6050_LightDMP_t dmp_state = {0};
static MPU6050_DMP_Config_t dmp_config = {0};
static MPU6050_Calibration_t calibration = {0};

/****************************************************************************************************************/
/*                                        核心算法实现                                                         */
/****************************************************************************************************************/

/**
 * @brief 快速平方根倒数算法（Quake算法优化版）
 * @param x: 输入值
 * @return 1/sqrt(x)的近似值
 */
float Fast_InvSqrt(float x)
{
    float halfx = 0.5f * x;
    float y = x;
    long i = *(long*)&y;
    i = 0x5f3759df - (i>>1);
    y = *(float*)&i;
    y = y * (1.5f - (halfx * y * y));
    return y;
}

/**
 * @brief Madgwick AHRS算法更新
 * @param gx, gy, gz: 陀螺仪数据 (弧度/秒)
 * @param ax, ay, az: 加速度计数据 (归一化)
 * @param mx, my, mz: 磁力计数据 (归一化)
 * @param use_mag: 是否使用磁力计数据
 */
void Madgwick_Update(float gx, float gy, float gz, float ax, float ay, float az, float mx, float my, float mz, bool use_mag)
{
    float recipNorm;
    float s0, s1, s2, s3;
    float qDot1, qDot2, qDot3, qDot4;
    float _2q0, _2q1, _2q2, _2q3, _4q0, _4q1, _4q2 ,_8q1, _8q2;
    float hx, hy, bx, bz;
    float _2q0mx, _2q0my, _2q0mz, _2q1mx, _2bx, _2bz, _4bx, _4bz;
    float _2q0q2, _2q2q3;
    float q0q0, q0q1, q0q2, q0q3, q1q1, q1q2, q1q3, q2q2, q2q3, q3q3;

    // 从状态结构获取当前四元数
    float q0 = dmp_state.q0;
    float q1 = dmp_state.q1;
    float q2 = dmp_state.q2;
    float q3 = dmp_state.q3;
    float beta = dmp_config.beta;

    // 四元数导数计算
    qDot1 = 0.5f * (-q1 * gx - q2 * gy - q3 * gz);
    qDot2 = 0.5f * (q0 * gx + q2 * gz - q3 * gy);
    qDot3 = 0.5f * (q0 * gy - q1 * gz + q3 * gx);
    qDot4 = 0.5f * (q0 * gz + q1 * gy - q2 * gx);

    // 如果加速度数据有效，计算重力参考方向的误差
    if(!((ax == 0.0f) && (ay == 0.0f) && (az == 0.0f))) {
        
        // 归一化加速度数据
        recipNorm = Fast_InvSqrt(ax * ax + ay * ay + az * az);
        ax *= recipNorm;
        ay *= recipNorm;
        az *= recipNorm;

        // 辅助变量以避免重复计算
        _2q0 = 2.0f * q0;
        _2q1 = 2.0f * q1;
        _2q2 = 2.0f * q2;
        _2q3 = 2.0f * q3;
        _4q0 = 4.0f * q0;
        _4q1 = 4.0f * q1;
        _4q2 = 4.0f * q2;
        _8q1 = 8.0f * q1;
        _8q2 = 8.0f * q2;
        q0q0 = q0 * q0;
        q1q1 = q1 * q1;
        q2q2 = q2 * q2;
        q3q3 = q3 * q3;

        // 梯度下降算法的目标函数
        s0 = _4q0 * q2q2 + _2q2 * ax + _4q0 * q1q1 - _2q1 * ay;
        s1 = _4q1 * q3q3 - _2q3 * ax + 4.0f * q0q0 * q1 - _2q0 * ay - _4q1 + _8q1 * q1q1 + _8q1 * q2q2 + _4q1 * az;
        s2 = 4.0f * q0q0 * q2 + _2q0 * ax + _4q2 * q3q3 - _2q3 * ay - _4q2 + _8q2 * q1q1 + _8q2 * q2q2 + _4q2 * az;
        s3 = 4.0f * q1q1 * q3 - _2q1 * ax + 4.0f * q2q2 * q3 - _2q2 * ay;
        
        // 归一化梯度
        recipNorm = Fast_InvSqrt(s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3);
        s0 *= recipNorm;
        s1 *= recipNorm;
        s2 *= recipNorm;
        s3 *= recipNorm;

        // 应用梯度反馈
        qDot1 -= beta * s0;
        qDot2 -= beta * s1;
        qDot3 -= beta * s2;
        qDot4 -= beta * s3;
    }

    // 如果磁力计数据有效且启用了磁力计，加入磁力计参考
    if(use_mag && !((mx == 0.0f) && (my == 0.0f) && (mz == 0.0f))) {
        // 归一化磁力计数据
        recipNorm = Fast_InvSqrt(mx * mx + my * my + mz * mz);
        mx *= recipNorm;
        my *= recipNorm;
        mz *= recipNorm;

        // 辅助变量
        _2q0mx = 2.0f * q0 * mx;
        _2q0my = 2.0f * q0 * my;
        _2q0mz = 2.0f * q0 * mz;
        _2q1mx = 2.0f * q1 * mx;
        // 这些变量已经在上面定义过，不需要重新定义
        //_2q0 = 2.0f * q0;
        //_2q1 = 2.0f * q1;
        //_2q2 = 2.0f * q2;
        //_2q3 = 2.0f * q3;
        _2q0q2 = 2.0f * q0 * q2;
        _2q2q3 = 2.0f * q2 * q3;
        // 这些变量已经在上面定义过，不需要重新定义
        //q0q0 = q0 * q0;
        q0q1 = q0 * q1;
        q0q2 = q0 * q2;
        q0q3 = q0 * q3;
        //q1q1 = q1 * q1;
        q1q2 = q1 * q2;
        q1q3 = q1 * q3;
        //q2q2 = q2 * q2;
        q2q3 = q2 * q3;
        //q3q3 = q3 * q3;

        // 参考地磁场方向
        hx = mx * (q0q0 - q1q1 - q2q2 + q3q3) + 2.0f * my * (q0q3 + q1q2) + 2.0f * mz * (q1q3 - q0q2);
        hy = 2.0f * mx * (q1q2 - q0q3) + my * (q0q0 - q1q1 + q2q2 - q3q3) + 2.0f * mz * (q0q1 + q2q3);
        bx = sqrtf(hx * hx + hy * hy);
        bz = 2.0f * mx * (q0q2 + q1q3) + 2.0f * my * (q2q3 - q0q1) + mz * (q0q0 - q1q1 - q2q2 + q3q3);

        // 磁力计梯度下降算法
        _2bx = 2.0f * bx;
        _2bz = 2.0f * bz;
        _4bx = 4.0f * bx;
        _4bz = 4.0f * bz;

        // 梯度下降算法校正步骤 - 磁力计
        s0 = -_2q2 * (2.0f * q1q3 - _2q0q2 - mx) + _2q1 * (2.0f * q0q1 + _2q2q3 - my) - _2bz * q2 * (_2bx * (0.5f - q2q2 - q3q3) + _2bz * (q1q3 - q0q2) - mx) + (-_2bx * q3 + _2bz * q1) * (_2bx * (q1q2 - q0q3) + _2bz * (q0q1 + q2q3) - my) + _2bx * q2 * (_2bx * (q0q2 + q1q3) + _2bz * (0.5f - q1q1 - q2q2) - mz);
        s1 = _2q3 * (2.0f * q1q3 - _2q0q2 - mx) + _2q0 * (2.0f * q0q1 + _2q2q3 - my) - 4.0f * q1 * (1 - 2.0f * q1q1 - 2.0f * q2q2 - mz) + _2bz * q3 * (_2bx * (0.5f - q2q2 - q3q3) + _2bz * (q1q3 - q0q2) - mx) + (_2bx * q2 + _2bz * q0) * (_2bx * (q1q2 - q0q3) + _2bz * (q0q1 + q2q3) - my) + (_2bx * q3 - _4bz * q1) * (_2bx * (q0q2 + q1q3) + _2bz * (0.5f - q1q1 - q2q2) - mz);
        s2 = -_2q0 * (2.0f * q1q3 - _2q0q2 - mx) + _2q3 * (2.0f * q0q1 + _2q2q3 - my) - 4.0f * q2 * (1 - 2.0f * q1q1 - 2.0f * q2q2 - mz) + (-_4bx * q2 - _2bz * q0) * (_2bx * (0.5f - q2q2 - q3q3) + _2bz * (q1q3 - q0q2) - mx) + (_2bx * q1 + _2bz * q3) * (_2bx * (q1q2 - q0q3) + _2bz * (q0q1 + q2q3) - my) + (_2bx * q0 - _4bz * q2) * (_2bx * (q0q2 + q1q3) + _2bz * (0.5f - q1q1 - q2q2) - mz);
        s3 = _2q1 * (2.0f * q1q3 - _2q0q2 - mx) + _2q2 * (2.0f * q0q1 + _2q2q3 - my) + (-_4bx * q3 + _2bz * q1) * (_2bx * (0.5f - q2q2 - q3q3) + _2bz * (q1q3 - q0q2) - mx) + (-_2bx * q0 + _2bz * q2) * (_2bx * (q1q2 - q0q3) + _2bz * (q0q1 + q2q3) - my) + _2bx * q1 * (_2bx * (q0q2 + q1q3) + _2bz * (0.5f - q1q1 - q2q2) - mz);

        // 归一化
        recipNorm = Fast_InvSqrt(s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3);
        s0 *= recipNorm;
        s1 *= recipNorm;
        s2 *= recipNorm;
        s3 *= recipNorm;

        // 应用磁力计梯度
        qDot1 -= beta * s0;
        qDot2 -= beta * s1;
        qDot3 -= beta * s2;
        qDot4 -= beta * s3;
    }

    // 积分四元数变化率得到四元数
    q0 += qDot1 * DMP_SAMPLE_PERIOD;
    q1 += qDot2 * DMP_SAMPLE_PERIOD;
    q2 += qDot3 * DMP_SAMPLE_PERIOD;
    q3 += qDot4 * DMP_SAMPLE_PERIOD;

    // 归一化四元数
    recipNorm = Fast_InvSqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3);
    q0 *= recipNorm;
    q1 *= recipNorm;
    q2 *= recipNorm;
    q3 *= recipNorm;

    // 更新状态
    dmp_state.q0 = q0;
    dmp_state.q1 = q1;
    dmp_state.q2 = q2;
    dmp_state.q3 = q3;
}

/**
 * @brief 四元数转欧拉角
 * @param q0, q1, q2, q3: 四元数
 * @param roll, pitch, yaw: 输出欧拉角(度)
 */
void Quaternion_To_Euler(float q0, float q1, float q2, float q3, float *roll, float *pitch, float *yaw)
{
    // Roll (x轴旋转)
    float sinr_cosp = 2.0f * (q0 * q1 + q2 * q3);
    float cosr_cosp = 1.0f - 2.0f * (q1 * q1 + q2 * q2);
    *roll = atan2f(sinr_cosp, cosr_cosp) * RAD_TO_DEG;

    // Pitch (y轴旋转)
    float sinp = 2.0f * (q0 * q2 - q3 * q1);
    if (fabsf(sinp) >= 1.0f) {
        *pitch = copysignf(90.0f, sinp); // 奇点处理
    } else {
        *pitch = asinf(sinp) * RAD_TO_DEG;
    }

    // Yaw (z轴旋转)
    float siny_cosp = 2.0f * (q0 * q3 + q1 * q2);
    float cosy_cosp = 1.0f - 2.0f * (q2 * q2 + q3 * q3);
    *yaw = atan2f(siny_cosp, cosy_cosp) * RAD_TO_DEG;
}

/****************************************************************************************************************/
/*                                        公共API实现                                                          */
/****************************************************************************************************************/

/**
 * @brief 初始化轻量级DMP
 * @param config: 配置参数，可以为NULL使用默认配置
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_Init(MPU6050_DMP_Config_t *config)
{
    // 先初始化基础MPU6050
    if (MPU6050_Init() != 0) {
        return 1;
    }

    // 设置默认配置
    if (config == NULL) {
        dmp_config.beta = MADGWICK_BETA;
        dmp_config.sample_rate = DMP_SAMPLE_RATE_HZ;
        dmp_config.auto_calibrate = true;
        dmp_config.calibrate_samples = 1000;
        dmp_config.use_magnetometer = true;  // 默认启用磁力计
    } else {
        dmp_config = *config;
    }

    // 配置传感器参数
    MPU6050_Set_Gyro_FSR(MPU_GYRO_FSR_2000DPS);
    MPU6050_Set_Accel_FSR(MPU_ACCEL_FSR_2G);
    MPU6050_Set_Rate(dmp_config.sample_rate);
    MPU6050_Set_LPF(dmp_config.sample_rate / 2);

    // 初始化四元数 (单位四元数)
    dmp_state.q0 = 1.0f;
    dmp_state.q1 = 0.0f;
    dmp_state.q2 = 0.0f;
    dmp_state.q3 = 0.0f;
    
    // 初始化其他状态
    dmp_state.initialized = false;
    dmp_state.update_count = 0;

    // 如果启用自动校准，执行校准
    if (dmp_config.auto_calibrate) {
        if (MPU6050_LightDMP_Calibrate(dmp_config.calibrate_samples) != 0) {
            return 1;
        }
    }

    dmp_state.initialized = true;
    return 0;
}

/**
 * @brief 传感器校准
 * @param samples: 校准采样数
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_Calibrate(uint16_t samples)
{
    MPU6050_Data_t sensor_data;
    float gyro_sum_x = 0, gyro_sum_y = 0, gyro_sum_z = 0;
    float accel_sum_x = 0, accel_sum_y = 0, accel_sum_z = 0;
    uint16_t valid_samples = 0;

    // 采集校准数据
    for (uint16_t i = 0; i < samples; i++) {
        if (MPU6050_Read_All_Data(&sensor_data) == 0) {
            gyro_sum_x += sensor_data.gyro_x;
            gyro_sum_y += sensor_data.gyro_y;
            gyro_sum_z += sensor_data.gyro_z;
            
            accel_sum_x += sensor_data.accel_x;
            accel_sum_y += sensor_data.accel_y;
            accel_sum_z += sensor_data.accel_z;
            
            valid_samples++;
        }
        mspm0_delay_ms(1); // 1ms延时
    }

    if (valid_samples < samples / 2) {
        return 1; // 校准失败，有效采样太少
    }

    // 计算偏移值
    calibration.gyro_offset_x = gyro_sum_x / valid_samples;
    calibration.gyro_offset_y = gyro_sum_y / valid_samples;
    calibration.gyro_offset_z = gyro_sum_z / valid_samples;
    
    calibration.accel_offset_x = accel_sum_x / valid_samples;
    calibration.accel_offset_y = accel_sum_y / valid_samples;
    calibration.accel_offset_z = (accel_sum_z / valid_samples) - ACCEL_SENSITIVITY_2G; // 减去重力

    calibration.calibrated = true;
    return 0;
}

/**
 * @brief 更新DMP数据（主循环调用）
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_Update(void)
{
    MPU6050_Data_t sensor_data;
    static HMC5883L_Data_t mag_data = {0};
    static bool mag_data_valid = false;
    
    // 定义指向外部变量的指针，如果外部变量不存在也不会导致链接错误
    bool *p_hmc5883l_ready = NULL;
    HMC5883L_Data_t *p_hmc_data = NULL;
    bool *p_hmc_data_valid = NULL;
    
    // 尝试获取外部符号地址（如果存在）
    #pragma weak hmc5883l_ready
    #pragma weak hmc_data
    #pragma weak hmc_data_valid
    extern bool hmc5883l_ready;
    extern HMC5883L_Data_t hmc_data;
    extern bool hmc_data_valid;
    
    // 设置指针（仅当外部符号存在时）
    if (&hmc5883l_ready != NULL) p_hmc5883l_ready = &hmc5883l_ready;
    if (&hmc_data != NULL) p_hmc_data = &hmc_data;
    if (&hmc_data_valid != NULL) p_hmc_data_valid = &hmc_data_valid;
    
    if (!dmp_state.initialized) {
        return 1;
    }

    // 读取传感器数据
    if (MPU6050_Read_All_Data(&sensor_data) != 0) {
        return 1;
    }

    // 应用校准偏移
    float gyro_x = ((float)(sensor_data.gyro_x - calibration.gyro_offset_x)) / GYRO_SENSITIVITY_2000DPS * DEG_TO_RAD;
    float gyro_y = ((float)(sensor_data.gyro_y - calibration.gyro_offset_y)) / GYRO_SENSITIVITY_2000DPS * DEG_TO_RAD;
    float gyro_z = ((float)(sensor_data.gyro_z - calibration.gyro_offset_z)) / GYRO_SENSITIVITY_2000DPS * DEG_TO_RAD;
    
    float accel_x = (float)(sensor_data.accel_x - calibration.accel_offset_x) / ACCEL_SENSITIVITY_2G;
    float accel_y = (float)(sensor_data.accel_y - calibration.accel_offset_y) / ACCEL_SENSITIVITY_2G;
    float accel_z = (float)(sensor_data.accel_z - calibration.accel_offset_z) / ACCEL_SENSITIVITY_2G;

    // 保存到状态结构（用于调试）
    dmp_state.gyro_x = gyro_x * RAD_TO_DEG;
    dmp_state.gyro_y = gyro_y * RAD_TO_DEG;
    dmp_state.gyro_z = gyro_z * RAD_TO_DEG;
    dmp_state.accel_x = accel_x;
    dmp_state.accel_y = accel_y;
    dmp_state.accel_z = accel_z;

    // 获取磁力计数据(如果有效)
    float mag_x = 0.0f, mag_y = 0.0f, mag_z = 0.0f;
    bool use_mag = dmp_config.use_magnetometer;
    
    // 尝试从外部源更新磁力计数据（如果可用）
    if (p_hmc5883l_ready != NULL && p_hmc_data_valid != NULL && p_hmc_data != NULL) {
        if (*p_hmc5883l_ready && *p_hmc_data_valid) {
            mag_data = *p_hmc_data;
            mag_data_valid = true;
        }
    }
    
    // 只有在磁力计数据有效时才使用它
    if (use_mag && mag_data_valid) {
        mag_x = mag_data.mag_x;
        mag_y = mag_data.mag_y;
        mag_z = mag_data.mag_z;
    } else {
        use_mag = false; // 如果没有有效数据，禁用磁力计
    }

    // 执行Madgwick算法更新
    Madgwick_Update(gyro_x, gyro_y, gyro_z, accel_x, accel_y, accel_z, mag_x, mag_y, mag_z, use_mag);

    // 计算欧拉角
    Quaternion_To_Euler(dmp_state.q0, dmp_state.q1, dmp_state.q2, dmp_state.q3,
                       &dmp_state.roll, &dmp_state.pitch, &dmp_state.yaw);

    dmp_state.update_count++;
    return 0;
}

/**
 * @brief 获取姿态角度
 * @param roll, pitch, yaw: 输出欧拉角(度)
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_GetAttitude(float *roll, float *pitch, float *yaw)
{
    if (!dmp_state.initialized || roll == NULL || pitch == NULL || yaw == NULL) {
        return 1;
    }

    *roll = dmp_state.roll;
    *pitch = dmp_state.pitch;
    *yaw = dmp_state.yaw;
    
    return 0;
}

/**
 * @brief 获取四元数
 * @param q0, q1, q2, q3: 输出四元数
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_GetQuaternion(float *q0, float *q1, float *q2, float *q3)
{
    if (!dmp_state.initialized || q0 == NULL || q1 == NULL || q2 == NULL || q3 == NULL) {
        return 1;
    }

    *q0 = dmp_state.q0;
    *q1 = dmp_state.q1;
    *q2 = dmp_state.q2;
    *q3 = dmp_state.q3;
    
    return 0;
}

/**
 * @brief 获取完整DMP数据
 * @param dmp_data: 输出数据结构
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_GetData(MPU6050_LightDMP_t *dmp_data)
{
    if (!dmp_state.initialized || dmp_data == NULL) {
        return 1;
    }

    *dmp_data = dmp_state;
    return 0;
}

/**
 * @brief 重置DMP状态
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_Reset(void)
{
    if (!dmp_state.initialized) {
        return 1;
    }

    // 重置四元数到单位四元数
    dmp_state.q0 = 1.0f;
    dmp_state.q1 = 0.0f;
    dmp_state.q2 = 0.0f;
    dmp_state.q3 = 0.0f;
    
    // 重置角度
    dmp_state.roll = 0.0f;
    dmp_state.pitch = 0.0f;
    dmp_state.yaw = 0.0f;
    
    dmp_state.update_count = 0;
    
    return 0;
}

/**
 * @brief 获取DMP状态
 * @param initialized: 是否已初始化
 * @param update_count: 更新计数
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_GetStatus(bool *initialized, uint32_t *update_count)
{
    if (initialized) *initialized = dmp_state.initialized;
    if (update_count) *update_count = dmp_state.update_count;
    return 0;
}

/**
 * @brief 设置DMP配置
 * @param config: 新配置
 * @return 0成功，1失败
 */
uint8_t MPU6050_LightDMP_SetConfig(MPU6050_DMP_Config_t *config)
{
    if (config == NULL) return 1;
    
    dmp_config = *config;
    
    // 如果已初始化，重新配置传感器
    if (dmp_state.initialized) {
        MPU6050_Set_Rate(dmp_config.sample_rate);
        MPU6050_Set_LPF(dmp_config.sample_rate / 2);
    }
    
    return 0;
} 